import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import datetime
from datetime import timedelta
from datetime import datetime
import time
"""
输入 53 * 53
"""


class SCN(nn.Module):
    def __init__(self):
        super(SCN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1, 1),  # (b, 32, 18, 18)
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1, 1),  # (b, 64, 9, 9)
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, 1, 1),  # (b, 128, 9, 9)
            nn.ReLU(),
            nn.Conv2d(128, 128, 3, 1, 1),  # (b, 128, 6, 6)
            nn.ReLU(),
            nn.Conv2d(128, 1, 3, 1, 1),  # (b, 2, 6, 6)
            # nn.AdaptiveAvgPool2d(1)
        )

    def forward(self, x):
        y = self.conv(x)
        # print(y.shape)
        return y.squeeze()


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.scn1 = SCN()
        # self.scn2 = SCN()

        # self.out = nn.AvgPool2d(6, 1)  # (b, 2, 1, 1)

    def forward(self, inputs):
        y = self.scn1(inputs)
        return y, y


if __name__ == '__main__':
    x = torch.randn(1, 1, 53, 53)
    # y = torch.randn(2, 2028)
    net = Net()
    # print(net)
    # print(y.detach().numpy())
    cur = datetime(2022, 8, 1, 0, 0)
    for i in range(100):
        # cur = datetime.now()
        y1, y2 = net(x)
        cur6 = cur + timedelta(minutes=6)
        cur12 = cur6 + timedelta(minutes=6)
        file1 = "{:04d}{:02d}{:02d}{:02d}{:02d}".format(cur.year, cur.month,
                                                         cur.day, cur.hour,
                                                         cur.minute)
        file2 = "{:04d}{:02d}{:02d}{:02d}{:02d}".format(cur6.year, cur6.month,
                                                         cur6.day, cur6.hour,
                                                         cur6.minute)
        file3 = "{:04d}{:02d}{:02d}{:02d}{:02d}".format(cur12.year, cur12.month,
                                                         cur12.day, cur12.hour,
                                                         cur12.minute)
        plt.imshow(x.squeeze())
        plt.savefig("E:\\ifstartjustcarryon\\RadarPred\\data0\\" + file1 + ".jpg", format='jpg')

        plt.imshow(y1.detach().numpy().squeeze())
        plt.savefig("E:\\ifstartjustcarryon\\RadarPred\\data6\\" + file2 + ".jpg", format='jpg')

        plt.imshow(y2.detach().numpy().squeeze())
        plt.savefig("E:\\ifstartjustcarryon\\RadarPred\\data12\\" + file3 + ".jpg", format='jpg')
        # print(file)
        # plt.show()
        plt.close()
        cur = cur + timedelta(minutes=6)
        time.sleep(5)
